#include <filesystem>
#include <iostream>
#include <atomic>
#include <future>
#include <iomanip>
#include <chrono>
#include "thread_pool.hpp"

namespace fs = std::filesystem;

// 工具函数：格式化文件大小显示
inline std::string formatSize(std::uintmax_t bytes)
{
    const char *units[] = {"B", "KB", "MB", "GB", "TB"};
    int unit = 0;
    double size = static_cast<double>(bytes);

    while (size >= 1024 && unit < 4)
    {
        size /= 1024;
        unit++;
    }

    char buffer[32];
    snprintf(buffer, sizeof(buffer), "%.2f %s", size, units[unit]);
    return std::string(buffer);
}

// 并行版本实现
class ParallelDirCalculator
{
public:
    ParallelDirCalculator(size_t thread_count = std::thread::hardware_concurrency())
        : pool_(thread_count), total_size_(0) {}

    std::uintmax_t calculateSize(const std::string &path)
    {
        if (!fs::exists(path))
        {
            std::cerr << "路径不存在: " << path << std::endl;
            return 0;
        }

        std::promise<void> completion_promise;
        auto completion_future = completion_promise.get_future();

        std::shared_ptr<std::atomic<int>> pending_tasks =
            std::make_shared<std::atomic<int>>(1);

        processDirectory(fs::path(path), pending_tasks, completion_promise);
        completion_future.wait();

        return total_size_.load();
    }

private:
    thread_pool pool_;
    std::atomic<std::uintmax_t> total_size_;

    void processDirectory(const fs::path &dir_path,
                          std::shared_ptr<std::atomic<int>> pending_tasks,
                          std::promise<void> &completion_promise)
    {

        pool_.add_task([this, dir_path, pending_tasks, &completion_promise]()
                       {
            try {
                for (const auto& entry : fs::directory_iterator(dir_path)) {
                    if (entry.is_regular_file()) {
                        total_size_ += entry.file_size();
                    }
                    else if (entry.is_directory()) {
                        pending_tasks->fetch_add(1);
                        processDirectory(entry.path(), pending_tasks, completion_promise);
                    }
                }
            }
            catch (const std::exception& e) {
                std::cerr << "处理目录时出错: " << e.what() << std::endl;
            }

            if (pending_tasks->fetch_sub(1) == 1) {
                completion_promise.set_value();
            } });
    }
};

// 串行版本实现
class SerialDirCalculator
{
public:
    std::uintmax_t calculateSize(const std::string &path)
    {
        if (!fs::exists(path))
        {
            std::cerr << "路径不存在: " << path << std::endl;
            return 0;
        }
        return processDirectory(fs::path(path));
    }

private:
    std::uintmax_t processDirectory(const fs::path &dir_path)
    {
        std::uintmax_t total_size = 0;

        try
        {
            for (const auto &entry : fs::directory_iterator(dir_path))
            {
                if (entry.is_regular_file())
                {
                    total_size += entry.file_size();
                }
                else if (entry.is_directory())
                {
                    total_size += processDirectory(entry.path());
                }
            }
        }
        catch (const std::exception &e)
        {
            std::cerr << "处理目录时出错: " << e.what() << std::endl;
        }

        return total_size;
    }
};

// 性能测试函数
void runComparisonTest(const std::string &path)
{
    std::cout << "\n=== 性能对比测试 ===" << std::endl;
    std::cout << "目录: " << path << std::endl;
    std::cout << "----------------------------------------" << std::endl;

    // 测试串行版本
    SerialDirCalculator serial_calculator;
    auto serial_start = std::chrono::high_resolution_clock::now();
    auto serial_size = serial_calculator.calculateSize(path);
    auto serial_end = std::chrono::high_resolution_clock::now();
    auto serial_duration = std::chrono::duration_cast<std::chrono::milliseconds>(serial_end - serial_start);

    // 测试并行版本
    ParallelDirCalculator parallel_calculator(std::thread::hardware_concurrency());
    auto parallel_start = std::chrono::high_resolution_clock::now();
    auto parallel_size = parallel_calculator.calculateSize(path);
    auto parallel_end = std::chrono::high_resolution_clock::now();
    auto parallel_duration = std::chrono::duration_cast<std::chrono::milliseconds>(parallel_end - parallel_start);

    // 输出结果
    std::cout << "串行版本:" << std::endl;
    std::cout << "总大小: " << formatSize(serial_size) << std::endl;
    std::cout << "耗时: " << serial_duration.count() << "ms" << std::endl;
    std::cout << "----------------------------------------" << std::endl;

    std::cout << "并行版本 (线程数: " << std::thread::hardware_concurrency() << "):" << std::endl;
    std::cout << "总大小: " << formatSize(parallel_size) << std::endl;
    std::cout << "耗时: " << parallel_duration.count() << "ms" << std::endl;
    std::cout << "----------------------------------------" << std::endl;

    // 计算性能提升
    double speedup = static_cast<double>(serial_duration.count()) /
                     parallel_duration.count();
    std::cout << "性能提升: " << std::fixed << std::setprecision(2)
              << speedup << "x" << std::endl;

    // 验证结果一致性
    if (serial_size != parallel_size)
    {
        std::cout << "警告：串行和并行计算结果不一致！" << std::endl;
        std::cout << "串行大小: " << serial_size << " bytes" << std::endl;
        std::cout << "并行大小: " << parallel_size << " bytes" << std::endl;
    }
}

int main()
{
    std::cout << "=== 目录大小计算性能对比测试 ===" << std::endl;
    std::cout << "测试环境：" << std::endl;
    std::cout << "CPU线程数: " << std::thread::hardware_concurrency() << std::endl;
    std::cout << "----------------------------------------" << std::endl;

    const char *test_paths[] = {
        "/usr/include",  // Linux系统头文件目录
        "/mnt/d/origin", // Windows系统目录
        ".",             // 当前目录
    };

    for (const auto &path : test_paths)
    {
        if (fs::exists(path))
        {
            runComparisonTest(path);
            std::cout << "\n";
        }
    }

    return 0;
}